#include <windows.h>
#include <iostream>
#include <vector>
#include <string>
#include <sstream>
#include <iomanip>
#include <cwchar>

// g++ main.cpp -o set_resolution.exe -Os -s -municode -static -ffunction-sections -fdata-sections -Wl,--gc-sections

// 设置控制台代码页为 UTF-8
void SetConsoleToUTF8() {
    SetConsoleOutputCP(CP_UTF8);
}

// 获取当前程序的文件名
std::wstring GetProgramName() {
    wchar_t path[MAX_PATH];
    GetModuleFileNameW(NULL, path, MAX_PATH); // 获取当前程序路径
    std::wstring fullPath(path);
    size_t lastSlash = fullPath.find_last_of(L"\\/"); // 找到最后一个斜杠
    if (lastSlash != std::wstring::npos) {
        return fullPath.substr(lastSlash + 1); // 提取文件名
    }
    return fullPath; // 如果没有斜杠，返回完整路径
}

// 结构体用于存储显示器信息
struct MonitorInfo {
    std::wstring deviceName;
    DEVMODEW currentMode; // 使用 DEVMODEW
};

// 回调函数，用于枚举显示器
BOOL CALLBACK MonitorEnumProc(HMONITOR hMonitor, HDC hdcMonitor, LPRECT lprcMonitor, LPARAM dwData) {
    std::vector<MonitorInfo>* monitors = reinterpret_cast<std::vector<MonitorInfo>*>(dwData);

    MONITORINFOEXW monitorInfo;
    monitorInfo.cbSize = sizeof(MONITORINFOEXW);
    if (GetMonitorInfoW(hMonitor, &monitorInfo)) {
        DEVMODEW devMode;
        devMode.dmSize = sizeof(DEVMODEW);
        devMode.dmDriverExtra = 0;

        // 获取当前显示设置
        if (EnumDisplaySettingsW(monitorInfo.szDevice, ENUM_CURRENT_SETTINGS, &devMode)) {
            MonitorInfo info;
            info.deviceName = monitorInfo.szDevice;
            info.currentMode = devMode; // 直接赋值 DEVMODEW
            monitors->push_back(info);
        }
    }
    return TRUE; // 继续枚举
}

// 设置显示器分辨率
bool SetDisplayResolution(const std::wstring& deviceName, int width, int height, int frequency) {
    DEVMODEW devMode;
    devMode.dmSize = sizeof(DEVMODEW);
    devMode.dmDriverExtra = 0;

    // 获取当前显示设置
    if (!EnumDisplaySettingsW(deviceName.c_str(), ENUM_CURRENT_SETTINGS, &devMode)) {
        std::wprintf(L"Failed to get current display settings for %s\n", deviceName.c_str());
        return false;
    }

    // 修改分辨率
    devMode.dmPelsWidth = width;
    devMode.dmPelsHeight = height;
    devMode.dmDisplayFrequency = frequency;
    devMode.dmFields = DM_PELSWIDTH | DM_PELSHEIGHT | DM_DISPLAYFREQUENCY;

    // 应用新的显示设置
    LONG result = ChangeDisplaySettingsExW(deviceName.c_str(), &devMode, NULL, CDS_UPDATEREGISTRY, NULL);
    switch (result) {
        case DISP_CHANGE_SUCCESSFUL:
            return true;
        case DISP_CHANGE_RESTART:
            std::wprintf(L"Resolution change requires restart for %s\n", deviceName.c_str());
            return false;
        default:
            std::wprintf(L"Failed to change resolution for %s\n", deviceName.c_str());
            return false;
    }
}

// 显示帮助信息
void ShowHelp() {
    std::wstring programName = GetProgramName(); // 获取当前程序名称
    std::wprintf(L"Usage: %s /x:WIDTH /y:HEIGHT /r:FREQUENCY /d:INDEX\n", programName.c_str());
    std::wprintf(L"Options:\n");
    std::wprintf(L"  /x:WIDTH      Set the width of the resolution (e.g., /x:1920)\n");
    std::wprintf(L"  /y:HEIGHT     Set the height of the resolution (e.g., /y:1080)\n");
    std::wprintf(L"  /r:FREQUENCY  Set the refresh rate (e.g., /r:60)\n");
    std::wprintf(L"  /d:INDEX      Specify the display index (e.g., /d:1)\n");
    std::wprintf(L"  /?            Show this help message\n");
}

// 解析命令行参数
bool ParseCommandLine(int argc, wchar_t* argv[], int& width, int& height, int& frequency, int& displayIndex) {
    width = 0;
    height = 0;
    frequency = 0;
    displayIndex = 1; // 默认设置为第一个显示器

    for (int i = 1; i < argc; i++) {
        std::wstring arg = argv[i];
        if (arg == L"/?") {
            ShowHelp();
            return false;
        } else if (arg.substr(0, 3) == L"/x:") {
            width = std::stoi(arg.substr(3));
        } else if (arg.substr(0, 3) == L"/y:") {
            height = std::stoi(arg.substr(3));
        } else if (arg.substr(0, 3) == L"/r:") {
            frequency = std::stoi(arg.substr(3));
        } else if (arg.substr(0, 3) == L"/d:") {
            displayIndex = std::stoi(arg.substr(3));
        } else {
            std::wprintf(L"Unknown argument: %s\n", arg.c_str());
            ShowHelp();
            return false;
        }
    }

    if (width <= 0 || height <= 0 || frequency <= 0 || displayIndex <= 0) {
        std::wprintf(L"Invalid arguments.\n");
        ShowHelp();
        return false;
    }

    return true;
}

int wmain(int argc, wchar_t* argv[]) {
    // 设置控制台代码页为 UTF-8
    SetConsoleToUTF8();

    // 如果没有参数，显示帮助信息
    if (argc == 1) {
        ShowHelp();
        return 0;
    }

    // 解析命令行参数
    int width, height, frequency, displayIndex;
    if (!ParseCommandLine(argc, argv, width, height, frequency, displayIndex)) {
        return 1;
    }

    // 枚举所有显示器
    std::vector<MonitorInfo> monitors;
    EnumDisplayMonitors(NULL, NULL, MonitorEnumProc, reinterpret_cast<LPARAM>(&monitors));

    // 打印显示器信息
    for (size_t monitorIndex = 0; monitorIndex < monitors.size(); monitorIndex++) {
        const MonitorInfo& info = monitors[monitorIndex];
        std::wprintf(L"Monitor %zu: %s\n", monitorIndex + 1, info.deviceName.c_str());
        std::wprintf(L"  Current resolution: %dx%d @ %dHz\n", info.currentMode.dmPelsWidth, info.currentMode.dmPelsHeight, info.currentMode.dmDisplayFrequency);
    }

    // 检查指定的显示器是否存在
    if (displayIndex > monitors.size()) {
        std::wprintf(L"Invalid display index. Only %zu monitors found.\n", monitors.size());
        return 1;
    }

    // 设置指定显示器的分辨率
    const std::wstring& deviceName = monitors[displayIndex - 1].deviceName;
    std::wprintf(L"Setting resolution for %s to %dx%d @ %dHz\n", deviceName.c_str(), width, height, frequency);
    if (SetDisplayResolution(deviceName, width, height, frequency)) {
        std::wprintf(L"Resolution set successfully!\n");
    }

    return 0;
}